# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""
import torch
from torch.utils.data import DataLoader
from torch.nn import functional as F
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.models import VGG11_Weights
from torchvision.transforms import ToTensor


def data_loader():
    print("data_loader")


def get_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    elif torch.backends.mps.is_available():
        return torch.device('mps')
    else:
        return "cpu"


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    print("")
    device = get_device()
    net: torch.nn.Module = torch.hub.load("pytorch/vision:v0.10.0", 'vgg11', weights=VGG11_Weights.DEFAULT)
    net.to(device)
    net.eval()

    data_test = CIFAR10(root="data", train=False, transform=ToTensor(), download=True)
    dataloader_test = DataLoader(data_test, batch_size=5)

    transform = transforms.Compose(
        [
            transforms.Resize(225),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ]
    )

    with open("../../assets/imagenet_classes.txt", "r") as f:
        categories = [s.strip() for s in f.readlines()]

    with torch.no_grad():
        for index, (data, label) in enumerate(dataloader_test):
            data, label = data.to(device), label.to(device)
            out = net(data)
            loss = F.cross_entropy(out, label)

            probabilities = F.softmax(out[0], dim=0)
            top5_prob, top5_catid = torch.topk(probabilities, 5)
            print("==================== loss:", loss)
            for i in range(top5_prob.size(0)):
                print(categories[top5_catid[i]], top5_prob[i].item())